/**
 * 与 CPU 有关的工具包
 * 
 * clear && gcc -o test_simd -g -Og -mavx -mavx2 -mavx512vl -DTEST_MAIN utils.c cpu_utils.c
 */

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <emmintrin.h> // SSE2
#include <immintrin.h> // AVX

#ifdef _WIN32
#elif (__linux__ || __unix__)
#include <unistd.h>
#endif

#include "utils.h"
#include "cpu_utils.h"
#include "cccompat.h"


std_err_t simd_int16_to_float(const int16_t *int16Array, float *floatArray, size_t size, size_t normalization_bits) {
    // 检查缓冲区起始地址的对齐情况
    // 对于 SSE 指令集，申请的空间要 16 字节对齐，对于 AVX 要 32 字节对齐
    if(!floatArray || !int16Array || !size || (((intptr_t)floatArray) & 31) || (((intptr_t)int16Array) & 31) || normalization_bits > 15) return STD_ERR_INVALID_DST_ADDR;
    float factor = (float)((1 << normalization_bits) - 1);    
    __m256 factor_vector = _mm256_set1_ps(factor); // 将归一化系数转换为浮点向量

    // 确保 size 是 8 的倍数
    if (size % 8 != 0) {
        // 处理剩余的元素
        for (size_t i = size - (size % 8); i < size; i++) {
            floatArray[i] = normalization_bits ? (float)(int16Array[i] / factor) : (float)int16Array[i];
        }
        size -= size % 8;
    }

    for (size_t i = 0; i < size; i += 8) {
        // Load 8 int16_t values into a 128-bit SIMD register
        __m128i int16_values = _mm_load_si128((__m128i*)(int16Array + i));

        // Convert int16_t to int32_t
        __m256i int32_values = _mm256_cvtepi16_epi32(int16_values);

        // Convert int32_t to float
        __m256 float_values = _mm256_cvtepi32_ps(int32_values);
        // 归一化
        if(normalization_bits) float_values = _mm256_div_ps(float_values, factor_vector);

        // Store the float values back into the float array
        _mm256_storeu_ps(floatArray + i, float_values);
    }

    return STD_ERR_OK;
}

std_err_t simd_float_to_int16(const float *floatArray, int16_t *int16Array, size_t size, size_t normalization_bits) {
    // 检查缓冲区起始地址的对齐情况
    // 对于 SSE 指令集，申请的空间要 16 字节对齐，对于 AVX 要 32 字节对齐
    if(!floatArray || !int16Array || !size || (((intptr_t)floatArray) & 31) || (((intptr_t)int16Array) & 31) || normalization_bits > 15) return STD_ERR_INVALID_DST_ADDR;
    float factor = (float)((1 << normalization_bits) - 1);
    __m256 factor_vector = _mm256_set1_ps(factor); // 将归一化系数转换为浮点向量

    // 确保 size 是 8 的倍数，因为 AVX 寄存器可以一次性处理 8 个 float 值
    if (size % 8 != 0) {
        // 处理剩余的元素
        for (size_t i = size - (size % 8); i < size; i++) {
            int16Array[i] = normalization_bits ? (int16_t)(floatArray[i] * factor) : (int16_t)floatArray[i];
        }
        size -= size % 8;
    }

    for (size_t i = 0; i < size; i += 8) {
        // 加载 8 个 float 值
        __m256 float_values = _mm256_load_ps(floatArray + i);
        // 反归一化
        if(normalization_bits) float_values = _mm256_mul_ps(float_values, factor_vector);

        // 将float转换为int32
        __m256i int32_values = _mm256_cvtps_epi32(float_values);

        // 将int32截断为int16
        __m128i int16_values_low = _mm256_cvtepi32_epi16(int32_values);

        // 存储结果
        _mm_store_si128((__m128i*)(int16Array + i), int16_values_low);
    }

    return STD_ERR_OK;
}

std_err_t int16_to_float(const int16_t *int16Array, float *floatArray, size_t size, size_t normalization_bits) {
    if(!int16Array || !floatArray || !size || normalization_bits > 15) return STD_ERR_INVALID_DST_ADDR;

    if(normalization_bits) {
        float factor = (float)((1 << normalization_bits) - 1);
        for (size_t i = 0; i < size; i++) floatArray[i] = (float)(int16Array[i] / factor);
    } else {
        for (size_t i = 0; i < size; i++) floatArray[i] = (float)(int16Array[i]);
    }

    return STD_ERR_OK;
}

std_err_t float_to_int16(const float *floatArray, int16_t *int16Array, size_t size, size_t normalization_bits) {
    if(!floatArray || !int16Array || !size || normalization_bits > 15) return STD_ERR_INVALID_DST_ADDR;

    if(normalization_bits) {
        float factor = (float)((1 << normalization_bits) - 1);
        for (size_t i = 0; i < size; i++) int16Array[i] = (int16_t)(floatArray[i] * factor);
    } else {
        for (size_t i = 0; i < size; i++) int16Array[i] = (int16_t)(floatArray[i]);
    }

    return STD_ERR_OK;
}

int simd_test(int argc, const char **argv) {
    size_t size = 320000; // 示例大小，必须是8的倍数
    int normalization = false;

    if(argc >= 2 && sscanf(argv[1], "%zu", &size));
    if(argc >= 3 && sscanf(argv[2], "%d", &normalization));
    time_us_t ts;

    int16_t *int16Array = (int16_t*)_mm_malloc(size * sizeof(int16_t), 32);
    float *floatArray = (float*)_mm_malloc(size * sizeof(float), 32);

    // 初始化int16Array
    for (size_t i = 0; i < size; i++) {
        int16Array[i] = i;
    }

    // loop 转换
    ts = time_us_now();
    int16_to_float(int16Array, floatArray, size, normalization);
    ts = time_us_elapsed(ts);
    printf("loop: %llu elapsed %llu us.\n", (FMTLLU)size, (FMTLLU)ts);
    memset(floatArray, 0, size * sizeof(float));

    // SIMD 转换
    ts = time_us_now();
    simd_int16_to_float(int16Array, floatArray, size, normalization);
    ts = time_us_elapsed(ts);
    printf("SIMD: %llu elapsed %llu us.\n", (FMTLLU)size, (FMTLLU)ts);

    // 打印结果
    for (size_t i = 0; i < 16; i++) {
        printf("floatArray[%zu] = %f\n", i, floatArray[i]);
    }

    // loop 转换
    ts = time_us_now();
    float_to_int16(floatArray, int16Array, size, normalization);
    ts = time_us_elapsed(ts);
    printf("loop: %llu elapsed %llu us.\n", (FMTLLU)size, (FMTLLU)ts);
    memset(int16Array, 0, size * sizeof(int16_t));

    // SIMD 转换
    ts = time_us_now();
    simd_float_to_int16(floatArray, int16Array, size, normalization);
    ts = time_us_elapsed(ts);
    printf("SIMD: %llu elapsed %llu us.\n", (FMTLLU)size, (FMTLLU)ts);

    // 打印结果
    for (size_t i = 0; i < 16; i++) {
        printf("int16Array[%zu] = %d\n", i, int16Array[i]);
    }

    // 释放内存
    _mm_free(int16Array);
    _mm_free(floatArray);

    return 0;
}

#ifdef TEST_MAIN
int main(int argc, const char **argv) {
    return simd_test(argc, argv);
}
#endif
